Add TSCATTER - #980
Conversation
7e63f3d to
e8965d1
Compare
| let extraClassDeclaration = [{ | ||
| bool hasIndexForm() { return static_cast<bool>(getIndexes()); } | ||
| bool hasMaskForm() { return static_cast<bool>(getMaskPatternAttr()); } | ||
| bool hasAxis() { return static_cast<bool>(getAxis()); } |
| PTODpsType:$src, | ||
| PTODpsType:$dst, | ||
| Optional<PTODpsType>:$indexes, | ||
| Optional<Index>:$axis, |
There was a problem hiding this comment.
axis应该和maskPattern作为attribute,不要作为operand,因为这是一个编译期静态信息
| %dst_mask = pto.alloc_tile : !pto.tile_buf<loc=vec, dtype=f16, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0> | ||
| pto.tscatter ins(%src_mask, {maskPattern = #pto.mask_pattern<P0101>} : !pto.tile_buf<loc=vec, dtype=f16, rows=1, cols=64, v_row=1, v_col=64, blayout=row_major, slayout=none_box, fractal=512, pad=0>) outs(%dst_mask : !pto.tile_buf<loc=vec, dtype=f16, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0>) | ||
| %axis = arith.constant 0 : index | ||
| pto.tscatter ins(%src_mask, {maskPattern = #pto.mask_pattern<P0101>} : !pto.tile_buf<loc=vec, dtype=f16, rows=1, cols=64, v_row=1, v_col=64, blayout=row_major, slayout=none_box, fractal=512, pad=0>, %axis : index) outs(%dst_mask : !pto.tile_buf<loc=vec, dtype=f16, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0>) |
There was a problem hiding this comment.
前面改了tscatter的verifier,要求dst的valid_col大于src的valid_col,这个例子是反过来的
| if (srcValid.size() != 2 || dstValid.size() != 2) | ||
| return emitOpError("expects src and dst to have rank-2 valid_shape"); | ||
|
|
||
| auto axis = getAxis(); |
There was a problem hiding this comment.
verifier需要拒绝同时提供index和axis
|
|
||
| def tscatter(src, dst, *, indexes=None, axis=None, mask_pattern=None): | ||
| """``pto.tscatter`` tile scatter wrapper.""" | ||
| _pto.tscatter( |
There was a problem hiding this comment.
DSL也要加校验,报错尽量提前。indexes和axis不能同时提供,axis的取值范围需要约束
| %src = pto.alloc_tile : !pto.tile_buf<loc=vec, dtype=f16, rows=1, cols=64, v_row=1, v_col=64, blayout=row_major, slayout=none_box, fractal=512, pad=0> | ||
| %dst = pto.alloc_tile : !pto.tile_buf<loc=vec, dtype=f16, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0> | ||
| pto.tscatter ins(%src, {maskPattern = #pto.mask_pattern<P0101>} : !pto.tile_buf<loc=vec, dtype=f16, rows=1, cols=64, v_row=1, v_col=64, blayout=row_major, slayout=none_box, fractal=512, pad=0>) outs(%dst : !pto.tile_buf<loc=vec, dtype=f16, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0>) | ||
| %src = pto.alloc_tile : !pto.tile_buf<loc=vec, dtype=f16, rows=1, cols=32, v_row=1, v_col=32, blayout=row_major, slayout=none_box, fractal=512, pad=0> |
There was a problem hiding this comment.
新增axis参数需要lit用例覆盖,确保ExpandTileOp和EmitC路径能正常工作,特别是EmitC路径下indexes form和axis form两套接口
| @tilelib.tile_template( | ||
| op="pto.tscatter", | ||
| target="a5", | ||
| name="template_tscatter_mask_p0101_col", | ||
| dtypes=_SCATTER_MASK_DTYPES, | ||
| iteration_axis="none", | ||
| op_engine="vector", | ||
| op_class="other", | ||
| layouts=("row_major",), | ||
| loop_depth=2, | ||
| is_post_update=False, | ||
| tags=("scatter", "mask", "p0101"), | ||
| constraints=(_axis_is_col,), | ||
| id=9, | ||
| ) | ||
| def template_tscatter_mask_p0101_col(src: pto.Tile, dst: pto.Tile, axis: pto.i32): | ||
| _template_tscatter_mask_col(src, dst, 1, 2) | ||
|
|
||
|
|
||
| @tilelib.tile_template( | ||
| op="pto.tscatter", | ||
| target="a5", | ||
| name="template_tscatter_mask_p1010_row", | ||
| dtypes=_SCATTER_MASK_DTYPES, | ||
| iteration_axis="none", | ||
| op_engine="vector", | ||
| op_class="other", | ||
| layouts=("row_major",), | ||
| loop_depth=2, | ||
| is_post_update=False, | ||
| tags=("scatter", "mask", "p1010"), | ||
| constraints=(_axis_is_row,), | ||
| id=3, | ||
| ) | ||
| def template_tscatter_mask_p1010_row(src: pto.Tile, dst: pto.Tile, axis: pto.i32): | ||
| _template_tscatter_mask_row(src, dst, (False,)) | ||
|
|
||
|
|
||
| @tilelib.tile_template( | ||
| op="pto.tscatter", | ||
| target="a5", | ||
| name="template_tscatter_mask_p1010_col", | ||
| dtypes=_SCATTER_MASK_DTYPES, | ||
| iteration_axis="none", | ||
| op_engine="vector", | ||
| op_class="other", | ||
| layouts=("row_major",), | ||
| loop_depth=2, | ||
| is_post_update=False, | ||
| tags=("scatter", "mask", "p1010"), | ||
| constraints=(_axis_is_col,), | ||
| id=10, | ||
| ) | ||
| def template_tscatter_mask_p1010_col(src: pto.Tile, dst: pto.Tile, axis: pto.i32): | ||
| _template_tscatter_mask_col(src, dst, 0, 2) |
There was a problem hiding this comment.
这些版本能合并吗?好像没必要分出这么多版本,TileOp的属性是可以在函数里面获取的
There was a problem hiding this comment.
index 模式的索引 dtype 规则三层不一致
C++ verifier(verifyIndexedForm):indexes 只能是 i16/i32 族,且有位宽规则——data 4B → idx 4B,data 2B → idx 2B,data 1B → idx 2B;
DSL 文档:"indexes must have i32 or ui32 dtype";
模板 dtype 矩阵:NUMERIC_DTYPES × {i32, ui32}
|
|
||
| from .metadata import ScalarSpec, VectorSpec, ViewSpec | ||
|
|
||
| _MASK_PATTERN_TOKENS = {"p0101", "p1010", "p0001", "p0010", "p0100", "p1000", "p1111"} |
There was a problem hiding this comment.
这个是TScatter专用的constraint,似乎没有必要加到公共constraint系统里,否则以后任何模板只要 tag 里碰巧含 "p0101" 之类就会被这个规则过滤。更干净的做法是把 "context attr 必须命中 template tags" 做成模板声明式的元数据(比如 require_tag_from=("mask_pattern",)),或者干脆用普通 constraint 函数表达。同理 _no_mask_pattern(mask_pattern=True) 这个"默认 True、收到字符串就为 False"的写法能工作但非常隐晦,依赖"约束参数名与 context attr 名巧合相同"这一隐式契约,建议至少加注释。
| def _make_index_inputs(name, src_dtype, idx_dtype, src_shape, idx_shape): | ||
| np_dt = npy_dtype(src_dtype) | ||
| np_idx_dt = npy_dtype(idx_dtype) | ||
| rng = np.random.RandomState(hash(name) & 0xFFFFFFFF) |
There was a problem hiding this comment.
建议改成zlib.crc32(name.encode()) 之类稳定种子
09c9fa3 to
e20ab52
Compare
| def _static_valid_dim(tile, index): | ||
| static_valid_shape = getattr(tile, "_template_static_valid_shape", None) | ||
| if static_valid_shape is None: | ||
| static_valid_shape = getattr(tile, "static_valid_shape", None) | ||
| if static_valid_shape is not None and static_valid_shape[index] is not None: | ||
| return static_valid_shape[index] | ||
| return tile.valid_shape[index] | ||
|
|
||
|
|
||
| def _template_tscatter_mask_row(src: pto.Tile, dst: pto.Tile, interleave_args): | ||
| dtype = dst.dtype | ||
| valid_rows = _static_valid_dim(src, 0) | ||
| valid_cols = _static_valid_dim(src, 1) |
There was a problem hiding this comment.
valid_row/valid_col可以通过src.valid_shape直接拿,不用自己实现接口
No description provided.